﻿using MathNet.Numerics;
using MathNet.Numerics.Providers.FourierTransform;
using MathNet.Numerics.Threading;
using System;
using System.Runtime.CompilerServices;
using System.Threading.Tasks;
using UtilZ.Dotnet.Ex.Model;

namespace UtilZ.Dotnet.Ex.FourierTransform
{
    /// <summary>
    /// Fourier变换类
    /// </summary>
    unsafe public partial class FourierEx
    {

        private static void PrimitiveForward(in SpanZ<Complex32> samples, FourierTransformScaling scaling)
        {
            if (samples.Length.IsPowerOfTwo())
            {
                if (samples.Length >= 1024)
                {
                    Radix2ForwardParallel(samples);
                }
                else
                {
                    Radix2Forward(samples);
                }
            }
            else
            {
                //BluesteinForward(samples);
                BluesteinConvolutionParallel(samples);
            }

            switch (scaling)
            {
                case FourierTransformScaling.SymmetricScaling:
                    HalfRescale(samples);
                    break;
                case FourierTransformScaling.ForwardScaling:
                    FullRescale(samples);
                    break;
            }
        }




        /// <summary>
        /// Fully rescale the FFT result.
        /// </summary>
        /// <param name="samples">Sample Vector.</param>
        private static void FullRescale(in SpanZ<Complex32> samples)
        {
            var scalingFactor = (float)1.0 / samples.Length;
            for (int i = 0; i < samples.Length; i++)
            {
                samples[i] *= scalingFactor;
            }
        }


        /// <summary>
        /// Half rescale the FFT result (e.g. for symmetric transforms).
        /// </summary>
        /// <param name="samples">Sample Vector.</param>
        private static void HalfRescale(in SpanZ<Complex32> samples)
        {
            var scalingFactor = (float)Math.Sqrt(1.0 / samples.Length);
            for (int i = 0; i < samples.Length; i++)
            {
                samples[i] *= scalingFactor;
            }
        }




        /// <summary>
        /// Convolution with the bluestein sequence (Parallel Version).
        /// </summary>
        /// <param name="samples">Sample Vector.</param>
        private static void BluesteinConvolutionParallel(in SpanZ<Complex32> samples)
        {
            int n = samples.Length;
            Complex32[] sequence = BluesteinSequence32(n);

            // Padding to power of two >= 2N–1 so we can apply Radix-2 FFT.
            int m = ((n << 1) - 1).CeilingToPowerOfTwo();
            //var b = new Complex32[m];
            //var a = new Complex32[m];
            var arrA = new Complex32[m];
            var arrB = new Complex32[m];
            Complex32* samplesPtr = samples.DataPtr;

            fixed (Complex32* ptrA = arrA, ptrB = arrB)
            {
                //SpanZ<Complex32> a = new SpanZ<Complex32>(ptrA, arrA.Length);
                //SpanZ<Complex32> b = new SpanZ<Complex32>(ptrB, arrB.Length);
                Complex32* a = ptrA, b = ptrB;

                CommonParallel.Invoke(
                    () =>
                    {
                        // Build and transform padded sequence b_k = exp(I*Pi*k^2/N)
                        for (int i = 0; i < n; i++)
                        {
                            b[i] = sequence[i];
                        }

                        for (int i = m - n + 1; i < m; i++)
                        {
                            b[i] = sequence[m - i];
                        }

                        Radix2Forward(new SpanZ<Complex32>(b, m));
                    },
                    () =>
                    {
                        // Build and transform padded sequence a_k = x_k * exp(-I*Pi*k^2/N)
                        for (int i = 0; i < n; i++)
                        {
                            a[i] = sequence[i].Conjugate() * samplesPtr[i];
                        }

                        Radix2Forward(new SpanZ<Complex32>(a, m));
                    });

                for (int i = 0; i < m; i++)
                {
                    a[i] *= b[i];
                }

                Radix2InverseParallel(new SpanZ<Complex32>(a, m));

                var nbinv = 1.0f / m;
                for (int i = 0; i < n; i++)
                {
                    samples[i] = nbinv * sequence[i].Conjugate() * a[i];
                }
            }
        }


        /// <summary>
        /// Radix-2 generic FFT for power-of-two sample vectors (Parallel Version).
        /// </summary>
        private static void Radix2InverseParallel(in SpanZ<Complex32> data)
        {
            Radix2Reorder(data);

            Complex32* ptr = data.DataPtr;
            //Complex32* ptr = data.DataPtr;
            int length = data.Length;
            for (var levelSize = 1; levelSize < length; levelSize *= 2)
            {
                var size = levelSize;

                CommonParallel.For(0, size, 64, (u, v) =>
                {
                    for (int i = u; i < v; i++)
                    {
                        Radix2Step(ptr, length, 1, size, i);
                    }
                });
            }
        }

        /// <summary>
        /// Generate the bluestein sequence for the provided problem size.
        /// </summary>
        /// <param name="n">Number of samples.</param>
        /// <returns>Bluestein sequence exp(I*Pi*k^2/N)</returns>
        private static Complex32[] BluesteinSequence32(int n)
        {
            double s = MathNet.Numerics.Constants.Pi / n;
            var sequence = new Complex32[n];

            // TODO: benchmark whether the second variation is significantly
            // faster than the former one. If not just use the former one always.
            if (n > BluesteinSequenceLengthThreshold)
            {
                Parallel.For(0, n, (k) =>
                {
                    double t = (s * k) * k;
                    sequence[k] = new Complex32((float)Math.Cos(t), (float)Math.Sin(t));
                });

                //for (int k = 0; k < sequence.Length; k++)
                //{
                //    double t = (s * k) * k;
                //    sequence[k] = new Complex32((float)Math.Cos(t), (float)Math.Sin(t));
                //}
            }
            else
            {
                Parallel.For(0, n, (k) =>
                {
                    double t = s * (k * k);
                    sequence[k] = new Complex32((float)Math.Cos(t), (float)Math.Sin(t));
                });

                //for (int k = 0; k < sequence.Length; k++)
                //{
                //    double t = s * (k * k);
                //    sequence[k] = new Complex32((float)Math.Cos(t), (float)Math.Sin(t));
                //}
            }

            return sequence;
        }











        /// <summary>
        /// Radix-2 generic FFT for power-of-two sized sample vectors.
        /// </summary>
        private static void Radix2Forward(in SpanZ<Complex32> data)
        {
            Radix2Reorder(data);
            Complex32* ptr = data.DataPtr;
            for (var levelSize = 1; levelSize < data.Length; levelSize *= 2)
            {
                for (var k = 0; k < levelSize; k++)
                {
                    Radix2Step(ptr, data.Length, -1, levelSize, k);
                }
            }
        }

        /// <summary>
        /// Radix-2 generic FFT for power-of-two sample vectors (Parallel Version).
        /// </summary>
        private static void Radix2ForwardParallel(in SpanZ<Complex32> data)
        {
            Radix2Reorder(data);

            //Complex32* ptr = data.DataPtr;
            Complex32* ptr = data.DataPtr;
            int length = data.Length;

            for (var levelSize = 1; levelSize < length; levelSize *= 2)
            {
                var size = levelSize;

                CommonParallel.For(0, size, 64, (u, v) =>
                {
                    for (int i = u; i < v; i++)
                    {
                        Radix2Step(ptr, length, -1, size, i);
                    }
                });
            }
        }

        /// <summary>
        /// Radix-2 Step Helper Method
        /// </summary>
        /// <param name="samples">Sample vector.</param>
        /// <param name="length">Sample vector length.</param>
        /// <param name="exponentSign">Fourier series exponent sign.</param>
        /// <param name="levelSize">Level Group Size.</param>
        /// <param name="k">Index inside of the level.</param>
#if !NET4_0
        [MethodImpl(MethodImplOptions.AggressiveInlining)]
#endif
        private static void Radix2Step(Complex32* samples, int length, int exponentSign, int levelSize, int k)
        {
            // Twiddle Factor
            var exponent = (exponentSign * k) * MathNet.Numerics.Constants.Pi / levelSize;
            var w = new Complex32((float)Math.Cos(exponent), (float)Math.Sin(exponent));

            var step = levelSize << 1;
            for (var i = k; i < length; i += step)
            {
                var ai = samples[i];
                var t = w * samples[i + levelSize];
                samples[i] = ai + t;
                samples[i + levelSize] = ai - t;
            }
        }



    }


}
